ステートマシーンを簡単に作るための「AWS Step Functions Data Science SDK」の紹介

クラスメソッド機械学習 on AWS Advent Calendar 2019

#サーバーレス

#機械学習

#Amazon SageMaker

#AWS Step Functions

#AWS

yoshim

2019.12.13

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

概要

当エントリは『機械学習 on AWS Advent Calendar 2019』の13日目のエントリです。

SageMakerを使ったワークフローをStepFunctionsで実装したい、となった時にとてもお世話になりそうな「AWS Step Functions Data Science SDK」のサンプルスクリプトを触ってみたので、簡単な解説エントリーです。
対象者は「Step FucntionsとSageMakerに興味がある人」です。

本エントリーの最終ゴールは、下記のようなステートマシーンを「AWS Step Functions Data Science SDK」を使って記述することです。

1.最初に

先日、AWSブログで紹介されていた「Amazon SageMaker 用 AWS Step Functions Data Science SDK」を使ってみます。

MLワークフローを実現する際に1番シンプルなのは「LambdaやFargateで必要な処理をまとめて実行すること」かと思います。ただ、これだと1スクリプトが担当する処理内容が大きくなりすぎて管理するときに大変なので、もう少し機能を分割して管理するために「StepFunctions」を使うこともあるかと思います。

「StepFunctions」も「cloud formation」や「aws sam」を使ってまとめてデプロイできると楽なのですが、分析やデータサイエンティストでない方が扱うにはまだ少し大変ですし、そこは本来やりたかったタスクではないはずです。
そこで、「AWS Step Functions Data Science SDK」を使って「SageMakerに付随するワークフローのステートマシーンの作成」で楽してハッピーになりましょう。

2.「AWS Step Functions Data Science SDK」について

「AWS Step Functions Data Science SDK」を使うと何ができるようになるのでしょう？
「1.最初に」でも少し言及した内容ですが、ソースリポジトリの「README.md」を確認しておきます。

Create steps that accomplish tasks. Chain those steps together into workflows. Include retry, succeed, or fail steps. Review a graphical representation and definition for your workflow. Create a workflow in AWS Step Functions. Start and review executions in AWS Step Functions.

参照:「AWS Step Functions Data Science SDK」のリポジトリ

つまり、ざっくり下記のような形ですね。

states(ステートマシーンの中での最小単位)を作って、繋げて、ステートマシーンを作成できる
retry、wait、pass、choice等のstateも作成可能
ステートマシーンの実行も可能

また、上記には特に言及されておりませんが、「既存のLambda関数をステートマシーン中で利用する」、「作ったステートマシーンのcfnテンプレート出力」といったこともできます。
StepFunctionsの記法に慣れていないうちは、かなり便利そうですね。

3.やること

とりあえずSageMakerのワークフローをステートマシーンで管理する部分の確認、ということで「こちら」のサンプルスクリプトを実行して内容を確認してみました。

最終的にできるステートマシーンは下記の通りです。
「モデルの学習」→「モデルの作成」→「バッチ変換JOBの実行」→「エンドポイント設定の作成」→「推論エンドポイントの作成」という、SageMakerを使う上でよくやる処理の流れです。

今回は上記のステートマシーンを「AWS Step Functions Data Science SDK」で作成し、ステートマシーンの情報をcfnテンプレートとして出力、といったところまでやってみます。

4.やってみる

基本的にはサンプルスクリプトの内容の通りですが、要点だけを抽出してまとめます。

4-1.各ステップを作成する

今回のステートマシーンで実行する、「モデルの学習」、「モデルの作成」、「バッチ変換JOBの実行」、「エンドポイント設定の作成」、「推論エンドポイントの作成」処理をそれぞれ定義します。
例えば、下記はモデルの学習についてのスクリプトです。

training_step = steps.TrainingStep(
    'Train Step', 
    estimator=xgb,
    data={
        'train': sagemaker.s3_input(train_s3_file, content_type='libsvm'),
        'validation': sagemaker.s3_input(validation_s3_file, content_type='libsvm')
    },
    job_name=execution_input['JobName']  
)

SageMakerに付随した処理については今回実施するもの以外に「ハイパーパラメータの自動チューニング」についても実装できます。詳細についてはこちらをご参照ください。

4-2.ステップを結合してワークフローにする

上記まででは、まだ各処理が独立して定義されているだけの状態ですので、「処理の順番」や「実行するIAMロール」等を指定した上で、ステートマシーンを作成します。

# ここまでの処理の順序を決定
workflow_definition = steps.Chain([
    training_step,
    model_step,
    transform_step,
    endpoint_config_step,
    endpoint_step
])

# 実行するIAMロールやステートマシーン名等を定義
workflow = Workflow(
    name='yoshim_MyTrainTransformDeploy_v1',
    definition=workflow_definition,
    role=workflow_execution_role,
    execution_input=execution_input
)

# ステートマシーンの作成
workflow.create()

これでステートマシーンが作成できました。

4-3.ステートマシーンの実行

続いて、実際にステートマシーンを実行してみましょう。
実行する時に必要な情報を定義して実行、と簡単に実行できます。
「4-1.各ステップを作成する」の工程でInputに何かを定義している場合はここで指定しましょう。

execution = workflow.execute(
    inputs={
        'JobName': 'yoshim-regression-{}'.format(uuid.uuid1().hex), # Each Sagemaker Job requires a unique name
        'ModelName': 'yoshim-regression-{}'.format(uuid.uuid1().hex), # Each Model requires a unique name,
        'EndpointName': 'yoshim-regression-{}'.format(uuid.uuid1().hex) # Each Endpoint requires a unique name,
    }
)

参照:stepfunctions.py

ちゃんと実行できました。

少し待つと、ステートマシーンが全て問題なく実行でき、

推論エンドポイントまで作成できており、今回実現したいことがちゃんとStepFunctionsで実現できていることが確認できました。

4-4.cfnテンプレートファイルの出力

折角ここまでできたので、最後にcfnテンプレートファイルの出力方法も確認しておきましょう。
今回実施したサンプルスクリプトには記載がなかったのですが、こちらを参考にしたところ下記のような形で出力できました。

from stepfunctions.workflow import cloudformation

cloudformation.build_cloudformation_template(workflow)

参照:cloudformation.py

出力結果は長いので割愛します。

5.まとめ

「AWS Step Functions Data Science SDK」でステートマシーンを作成することができました。
思ったより使いやすいですし、compute.pyを参照すると「Lambda」、「Glue」、「AWSBatch」、「ECS」の処理も追加できそうなので、実現できる範囲も広そうです。